package com.at.bigdata.spark.sql

import org.apache.spark.SparkConf
import org.apache.spark.sql.expressions.Aggregator
import org.apache.spark.sql.{Encoder, Encoders, Row, SparkSession, functions}

/**
 *
 * @author cdhuangchao3
 * @date 2023/5/27 8:05 PM
 */
object Spark03_SparkSql_UDAF1 {

  def main(args: Array[String]): Unit = {
    val sparkConf = new SparkConf().setMaster("local[*]").setAppName("operator")
    val spark = SparkSession.builder().config(sparkConf).getOrCreate()
    val df = spark.read.json("datas/user.json")
    df.createOrReplaceTempView("user")
    //    spark.sql("select age, 'Name:'+name from user").show()
    spark.udf.register("ageAvg", functions.udaf(new MyAvgUDAF))
    spark.sql("select ageAvg(age) from user").show()

    spark.close();
  }

  /**
   * 自定义平均值函数：计算年龄的平均值
   * 1、继承org.apache.spark.sql.expressions.Aggregator，定义泛型
   * IN：输入的数据类型
   * Buff：缓冲数据类型
   * OUT：输出数据类型
   *
   * 2、重写方法（6）
   */
  case class Buff(var total: Long, var count: Long)

  class MyAvgUDAF extends Aggregator[Long, Buff, Long] {
    // 初始值或0值
    override def zero: Buff = Buff(0L, 0L)

    // 根据输入的数据来更新缓冲区的数据
    override def reduce(b: Buff, a: Long): Buff = {
      b.total = b.total + a
      b.count = b.count + 1L
      b
    }

    // 合并缓冲区
    override def merge(b1: Buff, b2: Buff): Buff = {
      b1.total += b2.total
      b1.count += b2.count
      b1
    }

    // 计算
    override def finish(reduction: Buff): Long = reduction.total / reduction.count

    // 编码
    override def bufferEncoder: Encoder[Buff] = Encoders.product

    // 输出的编码操作
    override def outputEncoder: Encoder[Long] = Encoders.scalaLong
  }
}
